# 导入必要的库
import os
import warnings

# 设置环境变量以避免joblib在Windows上的CPU核心数检测问题
os.environ['LOKY_MAX_CPU_COUNT'] = '4'
# 设置TensorFlow环境变量以减少信息输出（如果有的话）
os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

# 抑制警告
warnings.filterwarnings('ignore', category=UserWarning)
warnings.filterwarnings('ignore', category=FutureWarning)

import gradio as gr
import numpy as np
import pickle
import cv2
from PIL import Image
import faiss

# 加载保存的模型
def load_models(model_dir='.'):
    """
    加载所有保存的集成模型
    """
    models = {}
    model_files = {
        '硬投票集成': 'voting_hard_voting_model.pkl',
        '软投票集成': 'voting_soft_voting_model.pkl',
        'KNN Bagging': 'bagging_knn_bagging_model.pkl',
        '决策树 Bagging': 'bagging_tree_bagging_model.pkl',
        'Stacking集成': 'stacking_stacking_model.pkl'
    }
    
    for name, filename in model_files.items():
        filepath = os.path.join(model_dir, filename)
        if os.path.exists(filepath):
            with open(filepath, 'rb') as f:
                models[name] = pickle.load(f)
            print(f"已加载模型: {name}")
        else:
            print(f"警告: 找不到模型文件 {filepath}")
    
    return models

# 提取图像特征（使用flat特征）
def extract_features(image):
    """
    从PIL图像提取flat特征
    
    参数:
    image: PIL Image对象
    
    返回:
    features: 特征向量 (1024维)
    """
    # 将PIL图像转换为numpy数组
    img_array = np.array(image)
    
    # 如果是RGBA，转换为RGB
    if len(img_array.shape) == 3 and img_array.shape[-1] == 4:
        img_array = cv2.cvtColor(img_array, cv2.COLOR_RGBA2RGB)
    
    # 转换为灰度图（因为训练时使用的是灰度图）
    if len(img_array.shape) == 3:
        img_gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
    else:
        img_gray = img_array
    
    # 调整大小为32x32
    img_resized = cv2.resize(img_gray, (32, 32))
    
    # 展平为1024维向量 (32*32=1024)
    features = img_resized.flatten().astype('float32')
    
    # L2归一化
    features = features.reshape(1, -1)
    faiss.normalize_L2(features)
    
    return features

# 预测函数（针对Stacking模型）
def predict_stacking(model_dict, features):
    """
    使用Stacking模型预测
    
    参数:
    model_dict: 包含'level1'和'meta'的字典
    features: 特征向量
    
    返回:
    prediction: 预测结果
    probabilities: 预测概率（如果可用）
    """
    level1_classifiers = model_dict['level1']
    meta_classifier = model_dict['meta']
    
    # 第一层预测
    meta_features = np.zeros((features.shape[0], len(level1_classifiers)))
    for i, clf in enumerate(level1_classifiers):
        meta_features[:, i] = clf.predict(features)
    
    # 第二层预测
    prediction = meta_classifier.predict(meta_features)[0]
    
    # 如果元分类器支持概率预测
    if hasattr(meta_classifier, 'predict_proba'):
        probabilities = meta_classifier.predict_proba(meta_features)[0]
        return prediction, probabilities
    else:
        return prediction, None

# 预测函数
def predict_image(image, model_name):
    """
    使用选定的模型预测图像
    
    参数:
    image: PIL Image对象
    model_name: 模型名称
    
    返回:
    result: 预测结果字符串
    """
    if image is None:
        return "⚠️ 请上传一张图片"
    
    if model_name not in models:
        return f"❌ 错误: 模型 '{model_name}' 未加载"
    
    try:
        # 提取特征
        features = extract_features(image)
        
        # 验证特征维度
        if features.shape[1] != 1024:
            return f"❌ 特征维度错误: {features.shape[1]}，应为1024"
        
        # 获取模型
        model = models[model_name]
        
        # 特殊处理Stacking模型
        if model_name == 'Stacking集成' and isinstance(model, dict):
            prediction, probabilities = predict_stacking(model, features)
        else:
            # 预测
            prediction = model.predict(features)[0]
            probabilities = None
            
            # 如果模型支持概率预测，获取置信度
            if hasattr(model, 'predict_proba'):
                probabilities = model.predict_proba(features)[0]
        
        # 转换预测结果
        class_name = "狗 🐕" if prediction == 1 else "猫 🐱"
        
        # 显示置信度
        if probabilities is not None:
            confidence = probabilities[int(prediction)] * 100
            result = f"✅ 预测结果: {class_name}\n📊 置信度: {confidence:.2f}%"
        else:
            result = f"✅ 预测结果: {class_name}"
        
        return result
    
    except Exception as e:
        import traceback
        error_msg = f"❌ 预测出错:\n{str(e)}\n\n详细信息:\n{traceback.format_exc()}"
        return error_msg

# 加载所有模型
print("=" * 50)
print("正在加载模型...")
print("=" * 50)
models = load_models('.')
print(f"\n成功加载 {len(models)} 个模型\n")

if not models:
    print("❌ 错误: 没有找到任何模型文件")
    print("请先运行: python train_ensemble.py 训练模型")
else:
    print(f"✅ 成功加载以下模型:")
    for name in models.keys():
        print(f"  - {name}")

# 创建Gradio界面
def create_interface():
    """
    创建Gradio Web界面
    """
    with gr.Blocks(title="猫狗分类器 - 集成学习") as demo:
        gr.Markdown("# 🐱🐕 猫狗图像分类器")
        gr.Markdown("上传一张猫或狗的图片，选择一个集成学习模型进行分类")
        
        with gr.Row():
            with gr.Column():
                # 图像输入
                image_input = gr.Image(
                    type="pil",
                    label="上传图片",
                    sources=["upload", "clipboard"]
                )
                
                # 模型选择
                model_dropdown = gr.Dropdown(
                    choices=list(models.keys()),
                    value=list(models.keys())[0] if models else None,
                    label="选择模型"
                )
                
                # 预测按钮
                predict_button = gr.Button("开始预测", variant="primary")
            
            with gr.Column():
                # 结果输出
                output_text = gr.Textbox(
                    label="预测结果",
                    lines=3,
                    interactive=False
                )
        
        # 添加示例说明
        gr.Markdown("### 模型说明")
        gr.Markdown("""
        - **硬投票集成**: 多个分类器投票，选择票数最多的类别
        - **软投票集成**: 多个分类器的概率平均，选择概率最高的类别
        - **KNN Bagging**: 基于KNN的装袋法集成
        - **决策树 Bagging**: 基于决策树的装袋法集成
        - **Stacking集成**: 两层分类器的堆叠集成
        """)
        
        # 绑定预测函数
        predict_button.click(
            fn=predict_image,
            inputs=[image_input, model_dropdown],
            outputs=output_text
        )
    
    return demo

# 启动应用
if __name__ == "__main__":
    if models:
        print("\n" + "=" * 50)
        print("正在启动Gradio Web应用...")
        print("=" * 50)
        demo = create_interface()
        print("\n🚀 启动Web服务器...")
        demo.launch(
            server_name="127.0.0.1",  # 只允许本地访问
            server_port=7860,         # 端口号
            share=False,              # 不创建公共链接
            inbrowser=True            # 自动打开浏览器
        )
    else:
        print("\n❌ 无法启动应用: 没有可用的模型")
        print("请先运行: python train_ensemble.py")
